__doc__ =   """ Classes and functions for computing and manipulating accessible surface areas.
            """

from operator import add, sub
from numpy import array, r_

from AbstractPropertyDict import AbstractAtomPropertyDict, AbstractResiduePropertyDict,\
                                 AbstractPropertyDict
from ZenExceptions import ZenExternalException
from ResidueProperties import AA_ASA
from Einput import einput
from Externals import guess_alt_loc
from Neighbors import Neighbors
from Externals import call_areaimol, call_stride
from Crunchers import asa_loop
from Math import sphere_points, coords_to_symmetry, coords_to_crystal

def get_surface(structure, symmetry_mode, crystal_mode, probe, points, bucket_size, MAXSYM):
    """ Calculates the atomic solvent-accessible surface area (ASA).
    """
    # this will not change in any run of run_asa
    # The default is to strip the whole structure from waters.
    spoints = sphere_points(points)  # point coordinates randomly distributed on the unit sphere.
    search_limit = 2 * (2.0 + probe)
    sh = structure.header

    # First we mask all water residues and chains with all residues masked (water chains).
    lattice_residues = einput(structure, 'R')
    lattice_residues._mask_children('H_HOH', 'eq', 'name')
    lattice_chains = einput(structure, 'C')
    lattice_chains._mask_children([], 'eq', 'values', function =True)

    # if not residues or chains are left not atoms to work with, abort with warning.
    if not lattice_chains.values():
        lattice_residues.unmask_children()
        lattice_chains.unmask_children()
        print 'Warning: No unmasked atoms to build lattice in chain(s)'
        return {}

    # these are all atoms we can work with
    lattice_atoms = einput(structure, 'A')
    lattice_atoms.dispatch('set_radius')

    # function which takes a selection of atoms (in the most common case all lattice_atoms
    # and lattice coords, which already were transformed to call the cython asa_loop properly
    # the reason for its existance is that biological symmetry sometimes operates only on a
    # selection of the asymmetric unit, so no common lattice_atoms choice is possible and further
    # sometimes multiple biological symmetry units are present (maybe partially) in one asymmetric
    # unit.

    def run_asa(atom_selection, lattice_coords, spoints, probe, bucket_size):
        # get array of radii of the selection of atoms e.g. one chain
        atom_radii  = array(atom_selection._data_children('radius')) + probe
        # ... and the corresponding coordinates
        atom_coords = array(atom_selection._data_children('coords', forgiving =False))
        # calculate bounding box in a form of an array
        atom_box = r_[atom_coords.min(axis=0) - search_limit, atom_coords.max(axis=0) + search_limit]
        # the lattice coordinates are atom coords after transformations in a 4D array
        # this array gets reshaped into an all_atoms x 3 array
        shape = lattice_coords.shape
        lattice_coords = lattice_coords.reshape((shape[0] * shape[1] * shape[2], shape[3]))
        # this calls the cython code which loops over all query atoms, surface points, and lattice atoms
        return asa_loop(atom_coords, lattice_coords, atom_radii, atom_radii, spoints, atom_box, probe, bucket_size)

    result = {}
    # the method arguments and crystal attributes are evaluated to call run_asa properly in the case of
    # structure.mode == "bio" multiple calls to run_asa could be required.
    if not symmetry_mode and not crystal_mode:
        lattice_coords = array(lattice_atoms._data_children('coords', forgiving =False))    # get all atoms
        lattice_coords = array([[lattice_coords]])                                          # fake 3D and 4D
        lattice_idx_to_id = dict(enumerate(lattice_atoms._data_children('get_full_id', forgiving =False, function =True)))
        asas = run_asa(lattice_atoms, lattice_coords, spoints, probe, bucket_size)
        for idx in xrange(asas.shape[0]):
            result[lattice_idx_to_id[idx]] = asas[idx]

    elif symmetry_mode in ('table', 'raw'):
        lattice_idx_to_id = dict(enumerate(lattice_atoms._data_children('get_full_id', forgiving =False, function =True)))
        lattice_coords = array(lattice_atoms._data_children('coords', forgiving =False))        # get all atoms
        # expand to symmetry, real 3D
        lattice_coords = coords_to_symmetry(lattice_coords, sh[symmetry_mode + '_fmx'], sh[symmetry_mode + '_omx'],\
                                                            sh[symmetry_mode + '_mxs'],    symmetry_mode)
        if crystal_mode:
            lattice_coords = coords_to_crystal(lattice_coords, sh[symmetry_mode + '_fmx'], sh[symmetry_mode + '_omx'], crystal_mode) # real 4D
        else:
            lattice_coords = array([lattice_coords])                                            # fake 4D
        asas = run_asa(lattice_atoms, lattice_coords, spoints, probe, bucket_size)
        for idx in xrange(asas.shape[0]):
            result[lattice_idx_to_id[idx]] = asas[idx]
    elif symmetry_mode == 'bio':
        start = 0
        for chain_ids, mx_num in sh['bio_cmx']:
            lattice_chains.mask_children()
            for chain_id in chain_ids:
                lattice_chains[structure.id + (0,) + chain_id].set_masked(False)
            select_atoms = einput(lattice_chains, 'A')
            select_idx_to_id = dict(enumerate(select_atoms._data_children('get_full_id', forgiving =False, function =True)))
            select_coords = array(select_atoms._data_children('coords', forgiving =False))        # get all atoms

            stop = start + mx_num
            select_coords = coords_to_symmetry(select_coords, sh['raw_fmx'], sh['raw_omx'], sh['bio_mxs'][start:stop], symmetry_mode)
            select_coords = array([select_coords])
            start = stop

            asas = run_asa(select_atoms, select_coords, spoints, probe, bucket_size)
            for idx in xrange(asas.shape[0]):
                result[select_idx_to_id[idx]] = asas[idx]

    lattice_residues.unmask_children()
    lattice_chains.unmask_children()
    return result


class AbstractSurface(AbstractAtomPropertyDict):
    """ Abstract class to map ASA values calculated onto atom entites.
        An atom's xtra dict. will contain data from the get_asa output.
    """
    def __init__(self, structure, xtra_key, asa_args):
        # these are some sane defaults
        asa_args['probe'] = (asa_args.get('probe') or 1.4)
        asa_args['points'] = (asa_args.get('points') or 960)
        asa_args['bucket_size'] = (asa_args.get('bucket_size') or 10)
        asa_args['MAXSYM'] =  (asa_args.get('MAXSYM') or 200000)
        asas = get_surface(structure, **asa_args)
        for query_id, asa in asas.iteritems():
            entity = structure.table['A'][query_id]
            entity.xtra[xtra_key] = asa
        AbstractAtomPropertyDict.__init__(self, asas)






class AbstractOperatorSurface(AbstractAtomPropertyDict):
    """ Abstract class to calculate and map the difference between two
        different ASA values for the same atom.
    """
    def __init__(self, structure, oper, xtra_key, xtra_keyA, xtra_keyB):
        structure.set_table()
        new_asas = {}
        for (full_id, atom) in structure.table['A'].iteritems():
            new_asa = eval(oper)((atom.xtra.get(xtra_keyA) or 0.0), (atom.xtra.get(xtra_keyB) or 0.0))
            atom.xtra[xtra_key] = new_asa
            new_asas[full_id] = new_asa
        AbstractAtomPropertyDict.__init__(self, new_asas)


class FreeSurface(AbstractAtomPropertyDict):
    def __init__(self, structure, xtra_key =None):
        xtra_key = (xtra_key or 'free_ASA')
        chains = einput(structure, 'C')
        free_asas = {}
        for chain in chains.itervalues(unmask =True):
            chains.mask_children()
            chain.masked = False
            free_asas.update(ASUSurface(structure, xtra_key))
        chains.unmask_children()
        AbstractAtomPropertyDict.__init__(self, free_asas)


class ASUSurface(AbstractSurface):
    def __init__(self, structure, xtra_key =None):
        xtra_key = (xtra_key or 'asu_ASA')
        asa_args = {'symmetry_mode':None,\
                    'crystal_mode':0}
        AbstractSurface.__init__(self, structure, xtra_key, asa_args)


class BioSurface(AbstractSurface):
    def __init__(self, structure, xtra_key =None):
        xtra_key = (xtra_key or 'bio_ASA')
        asa_args = {'symmetry_mode':'bio',\
                    'crystal_mode':0}
        AbstractSurface.__init__(self, structure, xtra_key, asa_args)


class UCSurface(AbstractSurface):
    def __init__(self, structure, xtra_key =None):
        xtra_key = (xtra_key or 'uc_ASA')
        asa_args = {'symmetry_mode':'raw',\
                    'crystal_mode':0}
        AbstractSurface.__init__(self, structure, xtra_key, asa_args)


class CrystalSurface(AbstractSurface):
    def __init__(self, structure, xtra_key =None):
        xtra_key = (xtra_key or 'crystal_ASA')
        asa_args = {'symmetry_mode':'raw',\
                    'crystal_mode':2}
        AbstractSurface.__init__(self, structure, xtra_key, asa_args)


class BioContactSurface(AbstractOperatorSurface):
    def __init__(self, structure, run =False, xtra_key =None):
        xtra_key = (xtra_key or 'bio_dASA')
        xtra_keyA = 'free_ASA'
        xtra_keyB = 'bio_ASA'
        if run:
            FreeSurface(structure)
            BioSurface(structure)
        AbstractOperatorSurface.__init__(self, structure, 'sub', xtra_key, xtra_keyA, xtra_keyB)


class TotalContactSurface(AbstractOperatorSurface):
    def __init__(self, structure, run =False, xtra_key =None):
        xtra_key = (xtra_key or 'total_dASA')
        xtra_keyA = 'free_ASA'
        xtra_keyB = 'crystal_ASA'
        if run:
            FreeSurface(structure)
            CrystalContactSurface(structure)
        AbstractOperatorSurface.__init__(self, structure, 'sub', xtra_key, xtra_keyA, xtra_keyB)


class CrystalContactSurface(AbstractOperatorSurface):
    def __init__(self, structure, run =False, xtra_key =None):
        xtra_key = (xtra_key or 'crystal_dASA')
        xtra_keyA = 'total_dASA'
        xtra_keyB = 'bio_dASA'
        if run:
            FreeSurface(structure)
            CrystalSurface(structure)
            BioSurface(structure)
            TotalContactSurface(structure, run =False)
            BioContactSurface(structure, run =False)
        AbstractOperatorSurface.__init__(self, structure, 'sub', xtra_key, xtra_keyA, xtra_keyB)


class ResidueSurface(AbstractResiduePropertyDict):
    def __init__(self, structure, xtra_key =None):
        xtra_key = (xtra_key or 'free_ASA')
        structure._data_propagate('sum', 'A', xtra_key, xtra =True)
        residues = einput(structure, 'R')
        asas = residues._propertydict_children(xtra_key, xtra =True)
        AbstractResiduePropertyDict.__init__(self, asas)


class RelativeSurface(AbstractResiduePropertyDict):
    """ Class to map atom surfaces onto parents.
    """
    def __init__(self, entities, srf_key =None):
        srf_key = (srf_key or 'free_ASA')
        surface_dict = {}
        residues = einput(entities,'R','residues')
        # xtra_key defines which atomic surfaces to work with
        srf_key_r = srf_key + '_R' # relative surface key
        for (id_, residue) in residues.iteritems():
            asa = residue.xtra.get(srf_key)
            try:
                r_asa = asa / AA_ASA[residue.get_name()]
                residue.xtra[srf_key_r] = r_asa
                surface_dict[id_] = (r_asa)
            except (KeyError, TypeError):
                pass
        AbstractResiduePropertyDict.__init__(self, surface_dict)


class StrideResidueASUSurface(AbstractResiduePropertyDict):
    """ Class to map ASA computed by Stride onto residues.
    """
    def __init__(self, pdb_file, structure, xtra_key =None, **kwargs):
        xtra_key = (xtra_key or 'stride_ASA')
        surfaces = call_stride(pdb_file, **kwargs)
        surface_dict = {}
        model = structure[(0,)]
        model.set_table(True)
        for ((chain_id, res_full_id), data) in surfaces.iteritems():
            try:
                residue = model[(chain_id,)][(res_full_id,)]
            except KeyError:
                print 'Residue %s %s %s cannot be found in: %s %s' % (res_full_id, model.get_full_id())
                continue
            surface_dict[residue.get_full_id()] = data[-1]  # data is a tuple
            residue.xtra[xtra_key] = data[-1]
        AbstractResiduePropertyDict.__init__(self, surface_dict)


class CoordinationNumberSurface(AbstractAtomPropertyDict, AbstractResiduePropertyDict):
    def __init__(self, ents, level =None, radius =12., ctype = None, xtra_key =None, **kwargs):
        xtra_key = (xtra_key or 'cn_SRF')
        level = (level or 'R')
        ctype = (ctype or 'coords')
        entities = einput(ents, level)
        neighbors = Neighbors(ents, level, ctype)
        surface_dict = {}
        for entity in entities:
            cn = len(neighbors.get_neighbors(entity, radius, level))
            surface_dict[entity.get_full_id()] = cn
            entity.xtra[xtra_key] = cn
        if level == 'A':
            AbstractAtomPropertyDict.__init__(self, surface_dict)
        elif level == 'R':
            AbstractResiduePropertyDict.__init__(self, surface_dict)
        else:
            AbstractPropertyDict.__init__(self, surface_dict)



#class AbstractAreaimolSurface(AbstractAtomPropertyDict):
    #""" Abstract class to run AreaIMol and map the results on atom entites.
        #An atoms's xtra dict will contain data from the pdb AreaIMol output.
    #"""
    #def __init__(self, pdb_file, model, xtra_key, keep =True, **kwargs):
        ## AreaIMol processes only the first model of the pdb file
        ## for zenPDB it does not make difference.
        #surface_dict = {}   # result

        #surfaces = call_areaimol(pdb_file, **kwargs)        # external call
        #model.set_table(True)                               # check this!
        #for (atom_id, atom_coord, atom_data) in surfaces:
            ## AreaIMol looses info about alt_loc and hetero flags
            ## we try to fix that here...
            #atom_full_id = guess_alt_loc(model, atom_id, atom_coord)
            #try:
                #atom = model.table['A'][atom_full_id]
            #except KeyError:
                #raise PDBExternalException, atom_full_id
            #if keep:
                #atom.xtra[xtra_key] = atom_data
            #surface_dict[atom_full_id] = atom_data
        #AbstractAtomPropertyDict.__init__(self, surface_dict)


#class AreaimolSurface(AbstractAreaimolSurface):
    #def __init__(self, pdb_file, model, xtra_key =None, **kwargs):
        #xtra_key = (xtra_key or 'EXP_AREAIMOL_ASA')
        ## Not all atoms will have 'EXP_AREAIMOL_ASA' assigned but almost .get('key')
        #AbstractAreaimolSurface.__init__(self, pdb_file, model, xtra_key, **kwargs)


#class ContactAreaimolSurface(AbstractAreaimolSurface):
    #def __init__(self, pdb_file, model, **kwargs):
        #kwargs['diffmode'] = 'IMOL'
        #kwargs['smode'] = 'IMOL'
        #kwargs['symmetry1'] = '1'
        #kwargs['symmetry2'] = model.get_parent().header['space_group'].replace(' ','')
        #kwargs['trans'] = '2'
        #xtra_key = 'CNT_AREAIMOL_ASA'
        #AbstractAreaimolSurface.__init__(self, pdb_file, model, xtra_key, **kwargs)


#EOF